from PySide6.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QLineEdit, QTextEdit, 
                             QProgressBar, QFileDialog, QTreeWidget, QTreeWidgetItem)
from PySide6.QtGui import QFont, QIcon, QPixmap
from PySide6.QtCore import Qt, QThread, Signal, QSize, QMutex, QMutexLocker
import sys
import os
import shutil
from PIL import Image  # 用于生成缩略图


# 生成缩略图的线程，一次处理一个文件夹
class ThumbnailGenerator(QThread):
    generated = Signal(str, str)  # 传递图片路径和缩略图路径
    finished = Signal(str)  # 完成信号，传递文件夹路径
    error = Signal(str, str)  # 错误信号，传递图片路径和错误信息

    def __init__(self, folder_path, image_extensions):
        super().__init__()
        self.folder_path = folder_path
        self.image_extensions = image_extensions
        self.running = True
        self.mutex = QMutex()
        self.temp_thumb_dir = os.path.join(os.path.dirname(__file__), "temp_thumbnails")
        os.makedirs(self.temp_thumb_dir, exist_ok=True)

    def stop(self):
        with QMutexLocker(self.mutex):
            self.running = False

    def run(self):
        try:
            for file in os.listdir(self.folder_path):
                # 检查是否需要停止
                with QMutexLocker(self.mutex):
                    if not self.running:
                        return

                file_path = os.path.join(self.folder_path, file)
                # 只处理图片文件
                if os.path.isfile(file_path) and file.lower().endswith(self.image_extensions):
                    try:
                        # 生成唯一的缩略图文件名
                        thumb_name = f"{hash(file_path)}.png"
                        thumb_path = os.path.join(self.temp_thumb_dir, thumb_name)
                        
                        # 如果缩略图已存在则跳过
                        if not os.path.exists(thumb_path):
                            with Image.open(file_path) as img:
                                img.thumbnail((64, 64))  # 生成64x64的缩略图
                                img.save(thumb_path)
                        
                        # 发送生成结果
                        self.generated.emit(file_path, thumb_path)
                    except Exception as e:
                        self.error.emit(file_path, str(e))

            self.finished.emit(self.folder_path)
        except Exception as e:
            self.error.emit(self.folder_path, str(e))


# 扫描文件夹结构的线程
class ScanThread(QThread):
    scan_progress = Signal(str)
    scan_finished = Signal(dict)  # 传递文件夹结构

    def __init__(self, folder_path):
        super().__init__()
        self.folder_path = folder_path
        self.image_extensions = ('.jpg', '.jpeg', '.png', '.webp')

    def run(self):
        # 构建文件夹结构：包含路径、类型、子节点和是否有图片的标记
        structure = {
            'type': 'folder', 
            'children': [], 
            'path': self.folder_path, 
            'has_image': False
        }
        self.build_structure(self.folder_path, structure)
        self.scan_finished.emit(structure)

    def build_structure(self, current_path, parent_node):
        try:
            for item in os.listdir(current_path):
                item_path = os.path.join(current_path, item)
                
                if os.path.isdir(item_path):
                    # 处理文件夹
                    folder_node = {
                        'type': 'folder',
                        'name': item,
                        'path': item_path,
                        'children': [],
                        'has_image': False
                    }
                    # 递归构建子文件夹结构
                    self.build_structure(item_path, folder_node)
                    
                    # 如果子文件夹有图片，标记当前文件夹也有图片
                    if folder_node['has_image']:
                        parent_node['has_image'] = True
                    
                    # 只有包含图片的文件夹才添加到结构中
                    if folder_node['has_image']:
                        parent_node['children'].append(folder_node)
                        self.scan_progress.emit(f"发现图片文件夹：{os.path.relpath(item_path, self.folder_path)}")
                
                elif item.lower().endswith(self.image_extensions):
                    # 处理图片文件
                    parent_node['has_image'] = True
                    file_node = {
                        'type': 'file',
                        'name': item,
                        'path': item_path
                    }
                    parent_node['children'].append(file_node)
                    self.scan_progress.emit(f"发现图片：{os.path.relpath(item_path, self.folder_path)}")
                    
        except Exception as e:
            self.scan_progress.emit(f"扫描错误：{str(e)}")


class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.image_list = []  # 存储所有图片路径
        self.folder_structure = None  # 文件夹结构
        self.thumbnail_generators = {}  # 正在运行的缩略图生成器
        self.image_extensions = ('.jpg', '.jpeg', '.png', '.webp')  # 支持的图片格式
        self.init_ui()

    def init_ui(self):
        # 窗口基本设置
        self.setWindowTitle("图片处理工具")
        self.setFixedSize(1000, 700)

        # 主布局
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(10)

        # 字体设置
        font = QFont()
        font.setPointSize(10)

        # 1. 选择图片文件夹区域
        folder_layout = QHBoxLayout()
        folder_label = QLabel("图片文件夹:")
        folder_label.setFont(font)
        self.folder_entry = QLineEdit()
        self.folder_entry.setFont(font)
        folder_btn = QPushButton("选择")
        folder_btn.setFont(font)
        folder_btn.clicked.connect(self.select_image_folder)
        folder_layout.addWidget(folder_label)
        folder_layout.addWidget(self.folder_entry)
        folder_layout.addWidget(folder_btn)

        # 2. 选择贴图文件区域
        sticker_layout = QHBoxLayout()
        sticker_label = QLabel("贴图文件:")
        sticker_label.setFont(font)
        self.sticker_entry = QLineEdit()
        self.sticker_entry.setFont(font)
        sticker_btn = QPushButton("选择")
        sticker_btn.setFont(font)
        sticker_btn.clicked.connect(self.select_sticker_path)
        sticker_layout.addWidget(sticker_label)
        sticker_layout.addWidget(self.sticker_entry)
        sticker_layout.addWidget(sticker_btn)

        # 3. 图片树状显示区域（带缩略图）
        tree_label = QLabel("待处理图片（带预览）:")
        tree_label.setFont(font)
        self.image_tree = QTreeWidget()
        self.image_tree.setFont(font)
        self.image_tree.setColumnCount(2)  # 两列：缩略图 + 文件名
        self.image_tree.setHeaderLabels(["预览", "文件名"])
        self.image_tree.setColumnWidth(0, 70)  # 缩略图列宽度
        self.image_tree.setAlternatingRowColors(True)  # 行交替颜色
        self.image_tree.setIconSize(QSize(64, 64))  # 图标大小
        self.image_tree.itemExpanded.connect(self.on_item_expanded)  # 监听展开事件

        # 4. 处理进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setAlignment(Qt.AlignHCenter)

        # 5. 开始处理按钮
        start_btn = QPushButton("开始处理")
        start_btn.setFont(font)
        start_btn.clicked.connect(self.start_processing)

        # 6. 日志信息框
        log_label = QLabel("处理日志:")
        log_label.setFont(font)
        self.info_text = QTextEdit()
        self.info_text.setReadOnly(True)
        self.info_text.setFont(font)

        # 添加所有组件到主布局
        main_layout.addLayout(folder_layout)
        main_layout.addLayout(sticker_layout)
        main_layout.addWidget(tree_label)
        main_layout.addWidget(self.image_tree)
        main_layout.addWidget(self.progress_bar)
        main_layout.addWidget(start_btn)
        main_layout.addWidget(log_label)
        main_layout.addWidget(self.info_text)

        self.setLayout(main_layout)

    def select_image_folder(self):
        """选择图片文件夹并开始扫描"""
        folder_path = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder_path:
            self.folder_entry.setText(folder_path)
            self.info_text.append(f"开始扫描文件夹结构: {folder_path}")
            
            # 清空之前的树状结构
            self.image_tree.clear()
            self.image_list = []
            
            # 停止所有正在运行的缩略图生成器
            for generator in self.thumbnail_generators.values():
                generator.stop()
            self.thumbnail_generators.clear()
            
            # 启动多线程扫描文件夹结构
            self.scan_thread = ScanThread(folder_path)
            self.scan_thread.scan_finished.connect(self.build_tree_structure)
            self.scan_thread.scan_progress.connect(self.update_scan_log)
            self.scan_thread.start()

    def select_sticker_path(self):
        """选择贴图文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择贴图文件", "", "Image files (*.jpg *.png *.jpeg *.webp)"
        )
        if file_path:
            self.sticker_entry.setText(file_path)
            self.info_text.append(f"已选择贴图文件: {os.path.basename(file_path)}")

    def update_scan_log(self, msg):
        """更新扫描日志"""
        self.info_text.append(msg)

    def build_tree_structure(self, structure):
        """构建树状结构"""
        self.folder_structure = structure
        # 创建根节点
        root_item = QTreeWidgetItem(["", os.path.basename(structure['path'])])
        root_item.setData(0, Qt.UserRole, structure['path'])  # 存储路径
        root_item.setData(1, Qt.UserRole, 'folder')  # 存储类型
        self.image_tree.addTopLevelItem(root_item)
        
        # 递归添加子节点
        self.add_tree_nodes(structure['children'], root_item)
        
        # 收集所有图片路径
        self.collect_all_images(structure)
        
        # 自动展开顶层节点
        root_item.setExpanded(True)
        self.info_text.append(f"扫描完成，共发现 {len(self.image_list)} 张图片")

    def add_tree_nodes(self, children, parent_item):
        """递归添加树节点"""
        for child in children:
            if child['type'] == 'folder':
                # 只添加有图片的文件夹
                item = QTreeWidgetItem(["", child['name']])
                item.setData(0, Qt.UserRole, child['path'])
                item.setData(1, Qt.UserRole, 'folder')
                parent_item.addChild(item)
                self.add_tree_nodes(child['children'], item)
            else:
                # 添加图片文件节点（初始无缩略图）
                item = QTreeWidgetItem(["", child['name']])
                item.setData(0, Qt.UserRole, child['path'])
                item.setData(1, Qt.UserRole, 'file')
                parent_item.addChild(item)

    def collect_all_images(self, node):
        """收集所有图片路径"""
        if node['type'] == 'folder':
            for child in node['children']:
                self.collect_all_images(child)
        else:
            self.image_list.append(node['path'])

    def on_item_expanded(self, item):
        """当节点展开时，生成该文件夹下图片的缩略图"""
        item_type = item.data(1, Qt.UserRole)
        if item_type != 'folder':
            return  # 只处理文件夹节点
        
        folder_path = item.data(0, Qt.UserRole)
        self.info_text.append(f"展开文件夹，开始生成缩略图: {os.path.basename(folder_path)}")
        
        # 如果已有生成器在运行，先停止
        if folder_path in self.thumbnail_generators:
            self.thumbnail_generators[folder_path].stop()
        
        # 创建并启动新的缩略图生成器
        generator = ThumbnailGenerator(folder_path, self.image_extensions)
        generator.generated.connect(self.update_thumbnail)
        generator.finished.connect(self.on_thumbnail_finished)
        generator.error.connect(self.on_thumbnail_error)
        self.thumbnail_generators[folder_path] = generator
        generator.start()

    def update_thumbnail(self, image_path, thumb_path):
        """更新图片节点的缩略图"""
        # 查找对应的树节点
        item = self.find_tree_item_by_path(image_path)
        if item and os.path.exists(thumb_path):
            pixmap = QPixmap(thumb_path)
            item.setIcon(0, QIcon(pixmap))

    def find_tree_item_by_path(self, target_path):
        """根据图片路径查找对应的树节点"""
        root = self.image_tree.invisibleRootItem()
        return self._recursive_find_item(root, target_path)

    def _recursive_find_item(self, parent_item, target_path):
        """递归查找树节点"""
        for i in range(parent_item.childCount()):
            child = parent_item.child(i)
            path = child.data(0, Qt.UserRole)
            if path == target_path:
                return child
            # 递归查找子节点
            found = self._recursive_find_item(child, target_path)
            if found:
                return found
        return None

    def on_thumbnail_finished(self, folder_path):
        """缩略图生成完成"""
        self.info_text.append(f"文件夹缩略图生成完成: {os.path.basename(folder_path)}")
        if folder_path in self.thumbnail_generators:
            del self.thumbnail_generators[folder_path]

    def on_thumbnail_error(self, path, error_msg):
        """处理缩略图生成错误"""
        self.info_text.append(f"生成缩略图错误 {os.path.basename(path)}: {error_msg}")

    def start_processing(self):
        """开始处理图片（后续可添加实际处理逻辑）"""
        if not self.image_list:
            self.info_text.append("请先选择图片文件夹并等待扫描完成")
            return
        if not self.sticker_entry.text():
            self.info_text.append("请先选择贴图文件")
            return
        self.info_text.append("开始处理图片...")

    def closeEvent(self, event):
        """窗口关闭时清理资源"""
        # 停止所有生成器
        for generator in self.thumbnail_generators.values():
            generator.stop()
        # 清理临时缩略图目录
        temp_thumb_dir = os.path.join(os.path.dirname(__file__), "temp_thumbnails")
        if os.path.exists(temp_thumb_dir):
            try:
                shutil.rmtree(temp_thumb_dir)
            except Exception as e:
                print(f"清理临时文件失败: {e}")
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
    